package main

import (
	"fmt"
	"sort"
)

func main() {
	arr := []int{1211, 212134, 172, 966, 111, 111}
	//BEGIN--->对arr数组离散化
    bak := make([]int, len(arr))
    copy(bak, arr)
    sort.Ints(bak) //排序
    unique(&bak) //去重
    Len:=len(arr)
    myTree := InitTree(Len) //构建主席树
    for i:=0; i<Len; i++ {
        //END--->离散化的数值插入到树中
		myTree.Insert(lower_bound(bak, arr[i]))  //二分查找arr[i] 在bak数组中的下标位置
	}
	fmt.Println(bak[myTree.Query(0, 5, 2)]) //查询arr[0-5]区间第2大是谁
    fmt.Println(bak[myTree.Query(1, 4, 1)]) //查询arr[1-4]区间第1大是谁
}

type Tree struct {	//主席树结构体
	root []*node	//用于存放每棵线段树的 "树根"
	index int		//构建树的时候用于根节点下标
}

type node struct {  //每棵线段树的结构
	sum, l, r int	//sum 值的计数, l, r 表示当前节点的左右区间
	left, right *node	//左右孩子节点
}

func InitTree(maxVal int) *Tree {	//初始化开辟主席树空间
	tree := &Tree{make([]*node, maxVal+1), 1}
	tree.root[0] = Init(maxVal)
	return tree
}
func Init(maxVal int) *node {	   //初始化第0个线段树的版本
	tree := new(node)
	tree.buildNode(0, maxVal)
	return tree
}
func (this *node) buildNode(le, ri int) {
	this.l, this.r = le, ri
	mid := (le + ri) >> 1
	if le == ri {return}
	this.left, this.right = new(node), new(node)
	this.left.buildNode(le, mid)
	this.right.buildNode(mid+1, ri)
}

func (this *Tree) Insert(x int) {       //插入数值
	this.root[this.index]=new(node)     //new一个新的线段树的树根
	insert(this.root[this.index-1], this.root[this.index], x)  //将前一个版本的树根和当前版本的树根传入
	this.index++
}
func insert(pre, cur *node, x int) {
	cur.l, cur.r, cur.sum = pre.l, pre.r, pre.sum+1             //只修改一条路径，这条路径上的值->，当前版本=历史版本+1
	if pre.l == pre.r {return}                               
	mid := (pre.l + pre.r) >> 1
	if x <= mid {                                               //当前值要插入到左区间
		cur.right,cur.left = pre.right, new(node)               //开辟左区间地址，递归左孩子
		insert(pre.left, cur.left, x)   
	} else {
		cur.left, cur.right = pre.left, new(node)               //同理递归右孩子
		insert(pre.right, cur.right, x)
	}
}

func (this *Tree) Query(l, r, k int) int { //查找区间第k大, 第k小同理
	return query(this.root[l], this.root[r+1], k)
}
func query(L, R *node, k int) int {
	if L.l == L.r {return L.l}
	tmp := R.right.sum - L.right.sum       //精华就在这里，利用前缀和原理， 如果k的值 小于 (R版本减L版本) 两个相减版本的右孩子sum值
	if k <= tmp {                          //第k大就在右区间继续递归
		return query(L.right, R.right, k)
	} else {
		return query(L.left, R.left, k-tmp) //否则进入左区间递归， 注意减去 两个相减版本的右孩子sum值，原理和权值线段树一样
	}
}

//离散化函数
//去重
func unique(arr *[]int) {
    slow, fast, Len := 0, 1, len(*arr)
    for fast < Len {
        if (*arr)[slow] != (*arr)[fast] {
            slow++
            (*arr)[slow] = (*arr)[fast]
        }
        fast++
    }
    *arr=(*arr)[:slow+1]
}
//二分  
func lower_bound(arr []int, x int) int {
        for left, right := 0, len(arr)-1; left <= right;  {
            mid := (right + left) >> 1
            if arr[mid] == x {
                return mid
            } else if arr[mid] > x {
                right = mid - 1
            } else {
                left = mid + 1
            }
        }
        return -1
}
//打印一颗线段树的所有节点,调试辅助用
func (this *node) ForPrint() {
	if this == nil {return}
	fmt.Println(this)
	this.left.ForPrint()
	this.right.ForPrint()
}